package com.leetcode.根据数据结构分类.前缀树;

import java.util.*;

/**
 * @author: xiaomi
 * @date: 2021/1/28
 * @description: 208. 实现 Trie (前缀树)
 * https://leetcode-cn.com/problems/implement-trie-prefix-tree/
 */
public class B_208_实现Trie {

    public static void main(String[] args) {
        Trie trie = new Trie();

        trie.insert("apple");
        // 返回 true
        boolean apple = trie.search("apple");
        System.out.println("apple = " + apple);
        // 返回 false
        boolean app = trie.search("app");
        System.out.println("app = " + app);
        // 返回 true
        boolean app1 = trie.startsWith("app");
        System.out.println("app2 = " + app1);
        trie.insert("app");
        // 返回 true
        System.out.println(trie.search("app"));

//        trie.insert("a");
//        System.out.println(trie.search("a"));
//        System.out.println(trie.startsWith("a"));

//        trie.insert("ab");
//        trie.search("abc");
//        trie.search("ab");
//        trie.startsWith("abc");
//        trie.startsWith("ab");
//        trie.insert("ab");
//        trie.search("abc");
//        trie.startsWith("abc");
//        trie.insert("abc");
//        System.out.println(trie.search("abc"));
//        System.out.println(trie.startsWith("abc"));

    }

    static class Trie {

        TrieNode first;

        /**
         * Initialize your data structure here.
         */
        public Trie() {
            first = new TrieNode();
        }

        /**
         * Inserts a word into the trie.
         */
        public void insert(String word) {
            //1.获取字符串的长度
            int len = word.length();
            TrieNode curNode = first;
            TrieNode preNode = curNode;
            int index = 0;
            while (curNode != null && index < len) {
                char c = word.charAt(index);
                preNode = curNode;
                if (!curNode.containsChar(c)) {
                    curNode = null;
                    //如果 Trie 中不包含 字符，就需要新建
                    break;
                }


                //此时子节点中包含了 c
                curNode = curNode.findNode(c);
                index++;
            }
            if (curNode != null && index != 0 && index == len - 1) {
                //如果已经遍历完 word
                return;
            }
            //此时 curNode == null || curNode 的子节点 不包含 c
            if (curNode == null) {
                curNode = preNode;
            }
            //将剩余的字符添加入 trie
            while (index < len) {
                TrieNode newNode = curNode.add(word.charAt(index++));
                curNode = newNode;
            }
            curNode.isEnd = true;
        }

        /**
         * Returns if the word is in the trie.
         * 必须是子节点
         */
        public boolean search(String word) {
            int len = word.length();
            TrieNode curNode = first;
            int index = 0;
            while (curNode != null && index < len) {
                char c = word.charAt(index);
                if (!curNode.containsChar(c)) {
                    //如果 Trie 中不包含 字符，就需要新建
                    return false;
                }

                //此时子节点中包含了 c
                curNode = curNode.findNode(c);
                index++;
            }
            if (index < len || curNode.isEnd == false) {
                //如果 word 中还有其他的长度没有比较 || 传参只是路径中的一部分而没有到达叶子
                //后者不能用叶子来判断，eg. app apple，所以考虑在 TrieNode 增加一个属性
                return false;
            }

            return true;
        }

        /**
         * Returns if there is any word in the trie that starts with the given prefix.
         */
        public boolean startsWith(String prefix) {
            int len = prefix.length();
            TrieNode curNode = first;
            int index = 0;
            while (curNode != null && index < len) {
                char c = prefix.charAt(index);
                if (!curNode.containsChar(c)) {
                    //如果 Trie 中不包含 字符，就需要新建
                    return false;
                }

                //此时子节点中包含了 c
                curNode = curNode.findNode(c);
                index++;
            }
            if (index < len) {
                //如果 word 中还有其他的长度没有比较
                return false;
            }

            return true;

        }
    }

    /**
     * 前缀树的节点
     */
    static class TrieNode {
        /**
         * 节点储存的元素,作为判断 equals 的唯一依据
         */
        Character c;
        /**
         * 子节点的数量
         */
        int size = 0;
        /**
         * 是否为 word 的末尾节点
         */
        boolean isEnd = false;
        /**
         * 节点的子节点集合
         * 用于遍历找到对应的节点
         */
        List<TrieNode> nodeList = new ArrayList<>();
        /**
         * 存储的字符副本(用于快速判断是否存在 字符)
         */
        Set<Character> characterSet = new HashSet<>();

        public TrieNode() {

        }

        public TrieNode(Character c) {
            this.c = c;
        }

        /**
         * 根据 char 添加子节点
         *
         * @param c
         * @return 新添加的子节点
         */
        public TrieNode add(char c) {
            TrieNode node = new TrieNode(c);
            add(node);
            return node;
        }

        /**
         * 根据新的节点，添加子节点
         *
         * @param node
         */
        public void add(TrieNode node) {
            this.nodeList.add(node);
            this.characterSet.add(node.c);
            size++;
        }

        public boolean containsChar(char c) {
            return characterSet.contains(c);
        }

        /**
         * 根据 字符 找到对应的子节点
         *
         * @param c
         * @return
         */
        public TrieNode findNode(char c) {
            for (int i = 0; i < size; i++) {
                if (nodeList.get(i).c.equals(c)) {
                    return nodeList.get(i);
                }
            }
            return null;
        }


        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            TrieNode trieNode = (TrieNode) o;
            return Objects.equals(c, trieNode.c);
        }

        @Override
        public int hashCode() {
            return Objects.hash(c);
        }
    }
}
